import argparse
import json
import os
import sys

sys.path.append(os.getcwd())

if __name__ == "__main__":
    if not os.path.exists("configs"):
        os.makedirs("configs")

    argparser = argparse.ArgumentParser()
    argparser.add_argument("--version", type=str, required=True, choices=["v1"])
    args = argparser.parse_args()

    yolo_config = {
        "MODEL": {
            "BACKBONE": "",
            "GRID_SIZE": 0,
            "BOUNDING_BOX_NUM": 0,
            "IMAGE_SIZE": 0,
        },
        "LOSS": {"LAMBDA_COORD": 0.0, "LAMBDA_NOOBJ": 0.0},
        "TRAIN": {
            "BATCH_SIZE": 0,
            "WORKERS_NUM": "auto",
            "EPOCHS": 0,
            "DEVICES": "auto",
        },
        "DATASET": {
            "SERIES": "",
            "VERSION": [],
            "ROOT_PATH": "",
            "AUTO_DOWNLOAD": True,
        },
    }

    if args.version == "v1":
        yolo_config["MODEL"]["GRID_SIZE"] = 7
        yolo_config["MODEL"]["BOUNDING_BOX_NUM"] = 2
        yolo_config["MODEL"]["IMAGE_SIZE"] = 448
        yolo_config["LOSS"]["LAMBDA_COORD"] = 5.0
        yolo_config["LOSS"]["LAMBDA_NOOBJ"] = 0.5

    with open(f"configs/yolo{args.version}.json", "w") as f:
        json.dump(yolo_config, f, indent=4)
